import torch
import torch.nn as nn
import torch.nn.functional as F
from ..module.conv import ConvModule
from ..module.init_weights import xavier_init
from ..module.transformer import TransformerBlock
from ..module.conv import DeepWise_PointWise_Conv
import matplotlib.pyplot as plt
import numpy as np
#将FPN最后一层的标准卷积替换为深度可分离卷积，并在其后使用Transformer
class FPN_Transformer(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 num_outs,
                 start_level=0,
                 end_level=-1,
                 conv_cfg=None,
                 norm_cfg=None,
                 activation=None
                 ):
        super(FPN_Transformer, self).__init__()
        assert isinstance(in_channels, list)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_ins = len(in_channels)
        self.num_outs = num_outs
        self.fp16_enabled = False

        if end_level == -1:
            self.backbone_end_level = self.num_ins
            assert num_outs >= self.num_ins - start_level
        else:
            # if end_level < inputs, no extra level is allowed
            self.backbone_end_level = end_level
            assert end_level <= len(in_channels)
            assert num_outs == end_level - start_level
        self.start_level = start_level
        self.end_level = end_level
        self.lateral_convs = nn.ModuleList()
        self.fpn_convs = nn.ModuleList()

        for i in range(self.start_level, self.backbone_end_level):
            l_conv = ConvModule(
                in_channels[i],
                out_channels,
                1,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                activation=activation,
                inplace=False)
            fpn_conv = DeepWise_PointWise_Conv(
                out_channels,
                out_channels,)

            self.lateral_convs.append(l_conv)
            self.fpn_convs.append(fpn_conv)
        self.transformer = TransformerBlock(out_channels, out_channels, 8,
                                            1, 4, 0.1, activation=activation)
        self.pos_embed = nn.Parameter(torch.zeros(400, 1, out_channels))
        # self.init_weights()

    # default init_weights for conv(msra) and norm in ConvModule
    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                xavier_init(m, distribution='uniform')

    def forward(self, inputs):
        assert len(inputs) == len(self.in_channels)

        # build laterals
        laterals = [
            lateral_conv(inputs[i + self.start_level])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]

        # build top-down path
        used_backbone_levels = len(laterals)
        for i in range(used_backbone_levels - 1, 0, -1):
            prev_shape = laterals[i - 1].shape[2:]
            laterals[i - 1] += F.interpolate(
                laterals[i], size=prev_shape, mode='bilinear')

        # build outputs
        outs = [
            self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
        ]

        feats = []
        gather_size = outs[1].shape[2:]

        # feats.append(F.adaptive_max_pool2d(outs[0],output_size=gather_size))
        # feats.append(outs[1])
        # feats.append(F.interpolate(outs[2], size=gather_size, mode='bilinear'))
        # bsf = sum(feats) / len(feats)



        # outs[0] = self.transformer(outs[0],self.pos_embed)
        # outs[1] = self.transformer(outs[1], self.pos_embed)
        # outs[2] = self.transformer(outs[3], self.pos_embed)


        # outs = [
        #     outs[0] + F.interpolate(bsf, size=outs[0].shape[2:], mode='bilinear'),
        #     outs[1] + bsf,
        #     outs[2] + F.interpolate(bsf, size=outs[2].shape[2:], mode='bilinear')
        # ]
        for i in range(3):
            print("shape is:")
            print(outs[i].shape)
        return tuple(outs)



